package prod_cons;

import java.util.Random;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;

import static prod_cons_pending.HeightCalculation.*;

public class ScalableTreeContainer<E> {
	public static class TreeNode<E> {
		public AtomicReference<E> data = new AtomicReference<E>();
		public AtomicInteger metadata = new AtomicInteger();
		public TreeNode<E> father;
		public int childPosition; // the position in the father's arr of
									// children (0 or 1)
		@SuppressWarnings("unchecked")
		public TreeNode<E>[] children = new TreeNode[2];

		public static int fatherIdx(int idx) {
			return idx >> 1;
		}

		public static int childPosition(int idx) { // 0 for left child, 1 for
													// right child
			return idx % 2;
		}

		public TreeNode(int idx, TreeNode<E> father) {
			this.father = father;
			this.childPosition = childPosition(idx);
		}
	}

	private static final int K = 3;
	private int initialHeight;
	
	private TreeNode<E>[][] tree; 
	public TreeNode<E> root;

	private static final ThreadLocal<Random> randomGenerator = new ThreadLocal<Random>() {
		protected Random initialValue() {
			return new Random();
		}
	};

	private static final ThreadLocal<Integer> lastUsedLevel = new ThreadLocal<Integer>() {
		protected Integer initialValue() {
			return new Integer(0);
		}
	};

	@SuppressWarnings("unchecked")
	public ScalableTreeContainer(int initialHeight) {
		this.initialHeight = initialHeight;
		tree = new TreeNode[initialHeight][];
		// create the root
		tree[0] = new TreeNode[1];
		tree[0][0] = new TreeNode<E>(0, null);
		root = tree[0][0];

		for (int level = 1; level < initialHeight; level++) {
			tree[level] = new TreeNode[1 << level]; // 2^level nodes for each
													// level
			for (int i = 0; i < tree[level].length; i++) {
				int fatherIdx = TreeNode.fatherIdx(i);
				tree[level][i] = new TreeNode<E>(i, tree[level - 1][fatherIdx]);
				tree[level - 1][fatherIdx].children[TreeNode.childPosition(i)] = tree[level][i];
			}
		}
	}

	public boolean put(E task) {
		int level = lastUsedLevel.get();
		int levelTrials = 1;
		
		Random random = randomGenerator.get();
		TreeNode<E> node = tree[level][random.nextInt(tree[level].length)];

		while (true) {
			if (level != 0 && !isDirty(node.father.metadata.get())) {
				// go to the father
				level = level - 1;
				levelTrials = 1;
				node = node.father;
				continue;
			}
			
			int meta = node.metadata.get();
			if (!isDirty(meta)) {
				// update dirty bit of the metadata
				// (tell to all the other putters that you're going to use this
				// node)
				int newMeta = setDirtyIncVersion(meta);
				if (node.metadata.compareAndSet(meta, newMeta)) {
					// insert the task
					node.data.set(task);
					updatePutMetadata(node);
					lastUsedLevel.set(level);
					return true;
				}
			}

			// the node has another task, go to the next node
			levelTrials--;
			if (levelTrials > 0) {
				// continue in the same level
				node = tree[level][random.nextInt(tree[level].length)];
			} else {
				// go down to the next level
				level++;
				if (level >= this.initialHeight) {
					lastUsedLevel.set(0);
					return false;
				}
				
				if (level == this.initialHeight - 1) {
					levelTrials = Math.min(level, K);
				} else {
					levelTrials = 1;
				}
				node = tree[level][random.nextInt(tree[level].length)];
			}
		}
	}

	public E get() {
		while (true) {
			if (!childrenTasks(root.metadata.get()) && root.data.get() == null)
				return null;
			E task = getAttempt();
			if (task != null) {
				return task;
			}
		}
	}

	public E getAttempt() {
		Random random = randomGenerator.get();
		TreeNode<E> node = root;
		while (true) {
			boolean hasData = (node.data.get() != null);
			int meta = node.metadata.get();
			if (!childrenTasks(meta) && hasData == false) {
				updateOnDemand(node);
				return null; // no tasks in this subtree
			}

			int nextChoice = nextChoice(random, hasData, meta);
			if (nextChoice == -1) {
				// first remove data, then update metadata
				E data = removeData(node);

				updateGetMetadata(node);
				return data;
			}
			node = node.children[nextChoice];
		}
	}

	private E removeData(TreeNode<E> node) {
		E data = node.data.get();
		if (data != null) {
			if (node.data.compareAndSet(data, null))
				return data;
		}
		return null;
	}
	
	private void addPending(AtomicInteger meta, int val) {
		while(true) {
			int metadata = meta.get();
			if (meta.compareAndSet(metadata, (val == 1) ? incPending(metadata) : decPending(metadata))) 
				return;
		}
	}

	private void updatePutMetadata(TreeNode<E> node) {
		// If the height has been changed go to the father
		// The height changes only if the task is added to the leaf node
		if (childrenTasks(node.metadata.get())) {
			// adding the task hasn't changed the height, you may return
			return;
		}
		updateNodeMetadata(node, true);
	}

	private void updateGetMetadata(TreeNode<E> node) {
		// If the height has been changed go to the father
		// The height changes only if the task is removed from the leaf node
		if (childrenTasks(node.metadata.get())) {
			// removing the task hasn't changed the height, you may return
			return;
		}
		updateNodeMetadata(node, false);
	}

	private void updateOnDemand(TreeNode<E> node) {
		TreeNode<E> father = node.father;
		if (father == null)
			return;
		while (true) {
			int oldFatherMeta = father.metadata.get();
			int meta = node.metadata.get();
			boolean hasTasks = ((node.data.get() != null) || leftTasks(meta) || rightTasks(meta));
			if (childTasks(oldFatherMeta, node.childPosition) == hasTasks)
				return;
			int newFatherMeta = childTasksIncVersion(oldFatherMeta,
					node.childPosition, (hasTasks == true) ? 1 : 0);
			if (father.metadata.compareAndSet(oldFatherMeta, newFatherMeta)) {
				updateNodeMetadata(father, false);
			}
		}
	}

	private void updateNodeMetadata(TreeNode<E> n, boolean val) {
		TreeNode<E> node = n;
		
		while (true) {
			TreeNode<E> father = node.father;
			if (father == null)
				return; // the father of the tree root is null

			int oldFatherMeta = father.metadata.get();
		
			int meta = node.metadata.get();
			boolean hasTasks = ((node.data.get() != null) || leftTasks(meta) || rightTasks(meta));
			if (hasTasks != val) {
				// my change has been eliminated, can return
				return;
			}
			boolean fatherVision = childTasks(oldFatherMeta, node.childPosition);
			
			// we want to update father's metadata in two following cases:
			// 1. father's metadata doesn't correspond to the current node's metadata
			// 2. father's metadata doesn't correspond to the previous metadata of the node 
			// (prio to the update that has happened in the previous iteration of while() loop)
			boolean needToUpdate = (fatherVision != hasTasks || getPending(meta) != 0);
			
			if (!needToUpdate) {
				node = father;
				continue;
			}
			
			AtomicInteger nodeMetadata = node.metadata;
			if (val == false) addPending(nodeMetadata, 1);
			// try to update father's metadata
			oldFatherMeta = father.metadata.get();
			int newFatherMeta = childTasksIncVersion(oldFatherMeta,	node.childPosition, (hasTasks == true) ? 1 : 0);
			if (father.metadata.compareAndSet(oldFatherMeta, newFatherMeta)) {
				node = father;
			} 
			if (val == false) addPending(nodeMetadata, -1);
		}
	}
}
